#!/usr/bin/python

# A simple ARM7TDMI assembler.
#
# Input Syntax:
# 	instructions:        <mnemonic> <cond> <flags> <operands>
#   labels:             <label_name>:
#   data definition:    <label_name>: <data_type> <data>
#			<data_type> ->  .ascii  = strings
#                       ->  .word   = 4-byte number
#						->  .hword  = 2-byte number
#						->  .byte   = 1-byte number/char
#
# Object code (output of this program):
#   -> the first instruction is a jump to the first instruction of the program
#   -> after this (byte 4 up), comes the data from the data definition block
#   -> then comes the main program code

# Version Log:
# v110705 - label / data definition support
# v110702 - basic skeleton

import sys
import struct

class Data_Def:
	def __init__(self, data_txt):
		if len(data_txt) < 3:
			print "Error: invalid data definition"
		self.label = data_txt[0]
		if data_txt[1] == ".ascii":
			self.type = "string"
			self.string = process_string(data_txt)
			self.size_bytes = len(self.string) + 1    # accounts for '\0'
		elif data_txt[1] == ".byte":
			self.type = "byte"
			self.data = chr(int(data_txt[2]))
			self.size_bytes = 1
		elif data_txt[1] == ".hword":
			self.type = "half_word"
			self.data = int(data_txt[2])
			self.size_bytes = 2
		elif data_txt[1] == ".word":
			self.type = "word"
			self.data = int(data_txt[2])
			self.size_bytes = 4
		else:
			print "unknown type: %s" % data_type
	def __str__(self):
		if self.type == "string":
			return "<type: " + self.type + "; data: " + self.string + ">"
		else:
			return "<type: " + self.type + "; data: " + self.data.__str__() + ">"
	def __repr__(self):
		return self.__repr__()

class Instruction:
	def __init__(self, inst_txt):
		self.txt = inst_txt
		self.size_bytes = 0
		self.prog_pos = -1
	def __str__(self):
		return "(" + self.txt.__str__() + "; bytes = " + self.size_bytes.__str__() + \
				"; position = " + self.prog_pos.__str__() + ")"
	def __repr__(self):
		return self.__str__()

#------------------------------------------------------------------------------

def process_string(inst):
	string = ""
	for i in inst[2:]:       # no label/directive
		string += i
	
	if string[0] != "\"":
		print "Error: invalid string format"
	
	jump = False             # ugly
	return_string = ""
	for i in xrange(1, len(string)):
		if jump == True:
			jump = False
			continue
		if string[i] == "\"":
			break
		elif i == "\\":
			return_string += get_escaped_char(string[i+1])
			jump = True
		else:
			return_string += string[i] 
	
	return return_string

#------------------------------------------------------------------------------

def get_escaped_char(char):
	if char == 'a':
		return '\a'
	elif char == 'b':
		return '\b'
	elif char == 'f':
		return '\f'
	elif char == 'n':
		return '\n'
	elif char == '\r':
		return '\r'
	elif char == 't':
		return '\t'
	elif char == 'v':
		return '\v'
	elif char == '\"':
		return '\"'
	elif char == '\\':
		return '\\'
	else:
		print "Error: unknown escape code"

#------------------------------------------------------------------------------
def cond_field(cond):
	if cond == 'eq':
		return 0x0
	elif cond == 'ne':
		return 0x1
	elif cond == 'cs':
		return 0x2
	elif cond == 'cc':
		return 0x3
	elif cond == 'mi':
		return 0x4
	elif cond == 'pl':
		return 0x5
	elif cond == 'vs':
		return 0x6
	elif cond == 'vc':
		return 0x7
	elif cond == 'hi':
		return 0x8
	elif cond == 'ls':
		return 0x9
	elif cond == 'ge':
		return 0xA
	elif cond == 'lt':
		return 0xB
	elif cond == 'gt':
		return 0xC
	elif cond == 'le':
		return 0xD
	elif cond == 'al':
		return 0xE
	else:
		print "unknown condition: %s" % cond
		
#------------------------------------------------------------------------------

def reg_field(reg):
	if reg[0] != 'r':
		print "error1: invalid register name: %s" % reg
	reg_num = -1
	
	try:
		reg_num = int(reg[1:])
	except ValueError:
		print "error2: invalid register name: %s" % reg
	
	return reg_num
	
#------------------------------------------------------------------------------

def print_little_endian(opcode, dest_file):
	byte_mask = 0xFF
	for i in xrange(4):
		dest_file.write("%c" % (opcode & byte_mask))
		opcode = opcode >> 8

#------------------------------------------------------------------------------

# BX instruction (branch and exchange)
# Format: bx cond r<n>
def emit_bx(instruction, dest_file):
	if len(instruction) != 3:
		print "error: invalid number of operands to bx"
	opcode = (cond_field(instruction[1]) << 28) | 0x12FFF10 | reg_field(instruction[2])
	print_little_endian(opcode, dest_file)

#------------------------------------------------------------------------------

# B/BL instruction (branch/branch with link)
# Format: bl cond offset
#  TODO: implement this instruction better
def emit_b(instruction, dest_file):
	if len(instruction) != 3:
		print "error: invalid number of operands to b/bl"
	link = 0x0B000000
	if instruction[0] == 'b':
		link = 0x0A000000
	opcode = (cond_field(instruction[1]) << 28) | link | int(instruction[2])
	print_little_endian(opcode, dest_file)

#------------------------------------------------------------------------------

def search_data(data, label):
	for i in data:
		if i.label == label:
			return i
	return None

#------------------------------------------------------------------------------

def assemble_code(instruction_list, dest_file):
	
	# Build Instruction objects list
	program = []
	for i in instruction_list:
		program.append(Instruction(i))
	
	# Build Data definition list
	data = []
	for i in instruction_list:
		if i[0][-1] == ':' and len(i) > 1 and i[1][0] == '.':  # data definition
			data.append(Data_Def(i))
		else:
			continue
	
	accum_addrs = 4   # first instruction is a jump to the first instruction of the program
	label_addrs = {}
	
	# Calculate size of data block and store label on dict
	for d in data:
		label_addrs[d.label] = accum_addrs
		accum_addrs += d.size_bytes
		if accum_addrs % 4 != 0:    # align data (not sure if it's necessary)
			accum_addrs += 4 - (accum_addrs % 4)
		

	# First pass: fill size field of instructions / create
	# dictionary with labels addresses
	sized_program = []
	for i in program:
		if i.txt[0][-1] != ':':   # labels have the form: <label_name>:
			#instruction -> 32 bits (4 bytes)
			i.size_bytes = 4
			i.prog_pos = accum_addrs
			accum_addrs = accum_addrs + 4
			sized_program.append(i)
		else:
			data_def = search_data(data, i.txt[0])
			if data_def != None:
				continue								
			label_addrs[i.txt[0]] = accum_addrs
	print sized_program		
	print label_addrs
"""
	for i in instruction_list:
		if i[0] == 'bx':
			emit_bx(i, dest_file)
		elif i[0] == 'b' or i[0] == 'bl':
			emit_b(i, dest_file)
		else:
			print "unknown instruction: %s" % i[0]
"""

#------------------------------------------------------------------------------

def main(args):
	
	# open source code file
	source_file = open(args[0], 'r')
	
	# open destination file
	dest_file = open(args[0] + ".o", 'wb')
	
	# read lines from source
	source_lines = source_file.readlines()
	print source_lines
	source_file.close()
	
	# trim lines
	trimmed_sl = []
	for i in source_lines:
		trimmed_sl.append(i.strip())
	print trimmed_sl
	
	# remove empty string/comments
	source_lines = []
	for i in trimmed_sl:
		if i != "" and i[0] != '@':
			source_lines.append(i)
	print source_lines
	
	# split instructions
	instructions = []
	for i in source_lines:
		instructions.append(i.split())
	print instructions
	
	# assemble and print codes to dest file
	assemble_code(instructions, dest_file)
	dest_file.close()
	
if __name__ == "__main__":
    main(sys.argv[1:])
